package com.gitee.feizns.quickstart.javacv;

import lombok.SneakyThrows;
import org.bytedeco.javacv.FFmpegFrameGrabber;
import org.bytedeco.javacv.FFmpegFrameRecorder;
import org.bytedeco.javacv.Frame;
import org.bytedeco.javacv.Java2DFrameConverter;

import java.awt.image.BufferedImage;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.List;

/**
 * java cv工具类
 * @author feizns
 * @since 2025/9/12
 */
public class JavaCv {

    /**
     * 提取关键帧
     * @param video 视频
     * @return {@link List }<{@link BufferedImage }>
     */
    @SneakyThrows
    public static List<BufferedImage> extractKeyFrames(InputStream video) {
        List<BufferedImage> frames = new ArrayList<>();
        FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(video);
        try(Java2DFrameConverter converter = new Java2DFrameConverter()) {
            grabber.start();
            Frame frame;
            // 按帧读取视频，只处理关键帧
            while ((frame = grabber.grabKeyFrame()) != null) {
                if (frame.image != null) {
                    // 转换为BufferedImage
                    frames.add(converter.convert(frame));
                }
            }
        } finally {
            grabber.stop();
            grabber.release();
        }
        return frames;
    }

    /**
     * 提取所有帧
     * @param video 视频
     * @return {@link List }<{@link BufferedImage }>
     */
    @SneakyThrows
    public static List<BufferedImage> extractFrames(InputStream video) {
        List<BufferedImage> frames = new ArrayList<>();
        FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(video);
        try(Java2DFrameConverter converter = new Java2DFrameConverter()) {
            grabber.start();
            Frame frame;
            // 按帧读取视频
            while ((frame = grabber.grabFrame()) != null) {
                if (frame.image != null) {
                    // 转换为BufferedImage
                    frames.add(converter.convert(frame));
                }
            }
        } finally {
            grabber.stop();
            grabber.release();
        }
        return frames;
    }

    /**
     * 提取帧
     * @param video 视频
     * @param interval 间隔（秒）
     * @return {@link List }<{@link BufferedImage }>
     */
    @SneakyThrows
    public static List<BufferedImage> extractFrames(InputStream video, int interval) {
        List<BufferedImage> frames = new ArrayList<>();
        FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(video);
        try(Java2DFrameConverter converter = new Java2DFrameConverter()) {
            grabber.start();
            // 获取视频帧率
            double frameRate = grabber.getVideoFrameRate();
            long duration = grabber.getLengthInVideoFrames();
            // 计算间隔帧数，interval单位为秒
            int frameInterval = (int) (frameRate * interval);
            // 每interval秒提取一帧
            for (int i = 0; i < duration; i += frameInterval) {
                grabber.setVideoFrameNumber(i);
                Frame frame = grabber.grabFrame();
                if (frame != null && frame.image != null) {
                    // 转换为BufferedImage
                    frames.add(converter.convert(frame));
                }
            }
        } finally {
            grabber.stop();
            grabber.release();
        }
        return frames;
    }

    /**
     * 提取音频
     * @param video 视频
     * @param audio 音频
     */
    public static void extractAudio(InputStream video, OutputStream audio) {
        extractAudio(video, audio, AudioType.mp3);
    }

    /**
     * 提取音频
     * @param video 视频
     * @param audio 音频
     * @param audioType 音频类型
     */
    @SneakyThrows
    public static void extractAudio(InputStream video, OutputStream audio, AudioType audioType) {
        FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(video);
        FFmpegFrameRecorder recorder = new FFmpegFrameRecorder(audio, grabber.getAudioChannels());
        // 0表示只有音频轨道
        try {
            grabber.start();
            recorder.setFormat(audioType.format());
            recorder.setAudioCodec(audioType.codec());
            recorder.setAudioBitrate(grabber.getAudioBitrate());
            recorder.setSampleRate(grabber.getSampleRate());
            recorder.setAudioChannels(grabber.getAudioChannels());
            recorder.start();
            Frame frame;
            while ((frame = grabber.grabFrame()) != null) {
                if (frame.samples != null) {
                    // 写入音频帧
                    recorder.record(frame);
                }
            }
        } finally {
            recorder.stop();
            recorder.release();
            recorder.close();
            grabber.stop();
            grabber.release();
        }
    }

}
